from __future__ import annotations
from pickletools import int4
import pandas as pd
import numpy as np
from typing import Union
from numpy import nan


#list 相关的四则运算及比较
class tdx(list):
    """重写运算符https://docs.python.org/zh-cn/3/library/operator.html"""
    #from __future__ import annotations 使得可以返回自身类型
    def __add__(self: list,other: Union[list,float,int]) -> tdx: #self:list没问题，因为isinstance(tdx,list)==True 返回tdx是为了返回值可以继续调用tdx方法
        """+"""
        if isinstance(other,(float,int)): #(float,int) 不可以换成Union[float,int]，Subscripted generics cannot be used with class and instance checks
            copy = self[:]
            for i in range(len(self)):
                copy[i]+=other  
            return tdx(copy)
        elif isinstance(other,list):
            if len(self)==len(other):
                copy = self[:]
                for i in range(len(other)):
                    copy[i]+=other[i]
                return tdx(copy)
            else:print("Type or Value Error")

    def __sub__(self: list,other: Union[list,float,int]) -> tdx:
        """-"""
        # if isinstance(self,list):
        if isinstance(other,(float,int)):
            copy = self[:]
            for i in range(len(self)):
                copy[i]-=other  
            return tdx(copy)
        elif isinstance(other,list):
            if len(self)==len(other):
                copy = self[:]
                for i in range(len(other)):
                    copy[i]-=other[i]
                return tdx(copy)
            else:print("Type or Value Error")

    def __mul__(self: list,other: Union[list,float,int]) -> tdx: 
        """*"""
        if isinstance(other,(float,int)):
            return tdx([i*other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i*j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __truediv__(self: list,other: Union[list,float,int]) -> tdx: 
        """/"""
        if isinstance(other,(float,int)):
            if other==0:
                raise ZeroDivisionError
            else:
                return tdx([i/other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i/j if j!=0 else nan for i,j in zip(self,other)])
        else:print("Type or Value Error")
    
    def __gt__(self: list,other: Union[list,float,int]) -> tdx: 
        """> greater than"""
        if isinstance(other,(float,int)):
            return tdx([i>other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i>j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __ge__(self: list,other: Union[list,float,int]) -> tdx: 
        """>= greater or equal than"""
        if isinstance(other,(float,int)):
            return tdx([i>=other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i>=j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __lt__(self: list,other: Union[list,float,int]) -> tdx: 
        """< litter than"""
        if isinstance(other,(float,int)):
            return tdx([i<other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i<j for i,j in zip(self,other)])
        else:print("Type or Value Error")
    
    def __le__(self: list,other: Union[list,float,int]) -> tdx: 
        """<= litter or equal than"""
        if isinstance(other,(float,int)):
            return tdx([i<=other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i<=j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __eq__ (self: list,other: Union[list,float,int]) -> tdx: 
        """== equal with"""
        if isinstance(other,(float,int)):
            return tdx([i==other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i==j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __ne__ (self: list,other: Union[list,float,int]) -> tdx: 
        """!= not equal with"""
        if isinstance(other,(float,int)):
            return tdx([i!=other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i!=j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __and__(self: list,other: Union[list,float,int]) -> tdx: 
        """and"""
        if isinstance(other,(float,int)):
            return tdx([i and other for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([i and j for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __abs__(self: list) -> tdx: 
        return tdx([abs(i) for i in self])

    def __max__(self:list,other: Union[list,float,int]) -> tdx: 
        if isinstance(other,(float,int)):
            return tdx([max(i,other) for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([max(i,j) for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __min__(self:list,other: Union[list,float,int]) -> tdx: 
        if isinstance(other,(float,int)):
            return tdx([min(i,other) for i in self])
        elif isinstance(other,list) and len(self)==len(other):
            return tdx([min(i,j) for i,j in zip(self,other)])
        else:print("Type or Value Error")

    def __if__(self:list,t_arg: Union[list,float,int],f_arg: Union[list,float,int]) -> tdx: 
        if isinstance(t_arg,(float,int)) and isinstance(f_arg,(float,int)):
            return tdx([t_arg if i else f_arg for i in self])
        elif isinstance(t_arg,list) and isinstance(f_arg,(float,int)) and len(self)==len(t_arg):
            return tdx([t_arg[i] if vlaue else f_arg for i,vlaue in enumerate(self)])
        elif isinstance(f_arg,list) and isinstance(t_arg,(float,int)) and len(self)==len(f_arg):
            return tdx([t_arg if vlaue else f_arg[i] for i,vlaue in enumerate(self)])
        elif isinstance(t_arg,list) and isinstance(f_arg,list) and len(self)==len(t_arg)==len(f_arg):
            return tdx([t_arg[i] if vlaue else f_arg[i] for i,vlaue in enumerate(self)])
        else:print("Type or Value Error")

    def __sma__(self:list,N:int,M:int)-> tdx:
        """
        #
        SMA(X，N，M) X的N日累积平均，M为权重
        如Y=(XM+Y(N-M))/N
        """
        sma_=[]
        for i in range(len(self)):
            if i==0:
                sma_.append(self[0])    
            elif np.isnan(self[i]) or np.isnan(sma_[-1]):
                sma_.append(self[i])
            else:
                sma_.append((M*self[i]+(N-M)*sma_[-1])/N)
        return tdx(sma_)

    def __ref__(self:list,window:int)-> tdx:
        return tdx(pd.Series(self).shift(window).to_list())
            
    def __ma__(self:list,window:int)->tdx:
        return tdx(pd.Series(self).rolling(window).mean().to_list())

    def __llv__(self:list,window:int)->tdx:
        return tdx(pd.Series(self).rolling(window).min().to_list())

    def __hhv__(self:list,window:int)->tdx:
        return tdx(pd.Series(self).rolling(window).max().to_list())

    def __ema__(self:list,N:int)->tdx:
        """
        若Yn=EMA(Xn，N)，则Yn=[2*Xn+(N-1)*Yn-1]/(N+1)，其中Yn-1表示上一周期的Y值。
        """
        ema_=[]
        for i in range(len(self)):
            if i==0:
                ema_.append(self[0])    
            elif np.isnan(self[i]) or np.isnan(ema_[-1]):
                ema_.append(self[i])
            else:
                ema_.append(2*self[i]/(N+1)+(N-1)/(N+1)*ema_[-1])
        return tdx(ema_)
        

    def __expma__(self:list,N:int)->tdx:#EMA指数平均数指标
        return tdx(self.__ema__(N))

    def __cross__(self:list,other:Union[list,float,int])->tdx: #上穿
        li=[]
        if isinstance(other,(float,int)):
            for i in range(len(self)):
                if i<len(self)-1:#防止报list_1[i+1]超出索引
                    li.append(self[i]<=other and self[i+1]>=other)
                else:
                    li.append(self[-1]==other)
            return tdx(li) 
        elif isinstance(other,list) and len(other)==len(self):
            for i in range(len(other)):
                if i<len(other)-1:#防止报list_1[i+1]超出索引
                    li.append(self[i]<=other[i] and self[i+1]>=other[i+1])
                else: #此时i已为len(list_)
                    li.append(self[i]==other[i])
            return tdx(li) 
        else:
            print("Type or Value Error")
        
    def __downcross__(self:list,other:Union[list,float,int])->tdx: #下穿
        li=[]
        if isinstance(other,(float,int)):
            for i in range(len(self)):
                if i<len(self)-1:#防止报list_1[i+1]超出索引
                    li.append(self[i]>=other and self[i+1]<=other)
                else:
                    li.append(self[-1]==other)
            return tdx(li) 
        elif isinstance(other,list) and len(other)==len(self):
            for i in range(len(other)):
                if i<len(other)-1:#防止报list_1[i+1]超出索引
                    li.append(self[i]>=other[i] and self[i+1]<=other[i+1])
                else: #此时i已为len(list_)
                    li.append(self[i]==other[i])
            return tdx(li) 
        else:
            print("Type or Value Error")

    def __std__(self:list,N:int)->tdx:             #求序列的N日标准差，返回序列    
        return  tdx(pd.Series(self).rolling(N).std(ddof=0).to_list())

    def __tfilter__(BUY:list,SELL:list,N:int4)->tdx:  #用法:TFILTER(买入条件,卖出条件,N);过滤掉买入(卖出)信号发出后,下一个反向信号发出前的所有买入(卖出)信号,
                            # N=1表示仅对买入信号过滤;
                            # N=2表示仅对卖出信号过滤;
                            # N=0表示对买入和卖出信号都过滤,返回1,2表示买入或卖出条件成立;
                            # 同一K线上只能有一个信号;
        res=[]
        #先找到买信号，再找到卖信号，如果重复，买信号为0，无效，买卖信号中间出现的买信号为0
        for i in range(len(BUY)):
            if BUY[i] and SELL[i]:
                BUY[i]=0
                SELL[i]=0
        if N==1:#buy filter
            buy_index=0
            while(buy_index<len(BUY)):
                buy=BUY[buy_index]
                if buy:
                    res.append(1)
                    i=buy_index
                    while((not SELL[i]) and i<len(SELL)-1):
                        res.append(0)
                        i+=1
                    buy_index=i+1
                else:
                    buy_index+=1
                    res.append(0)
        if N==2:#sell filter
            sell_index=0
            while(sell_index<len(SELL)):
                sell=SELL[sell_index]
                if sell:
                    res.append(1)
                    i=sell_index
                    while((not BUY[i]) and i<len(SELL)-1):
                        res.append(0)
                        i+=1
                    sell_index=i+1
                else:
                    sell_index+=1
                    res.append(0)
        if N==0: #返回1/2表示买入/卖出条件成立
            res_buy=[]
            buy_index=0
            while(buy_index<len(BUY)):
                buy=BUY[buy_index]
                if buy:
                    res_buy.append(1)
                    i=buy_index
                    while((not SELL[i]) and i<len(SELL)-1):
                        res_buy.append(0)
                        i+=1
                    buy_index=i+1
                else:
                    buy_index+=1
                    res_buy.append(0)

            res_sell=[]
            sell_index=0
            while(sell_index<len(SELL)):
                sell=SELL[sell_index]
                if sell:
                    res_sell.append(2)
                    i=sell_index
                    while((not BUY[i]) and i<len(SELL)-1):
                        res_sell.append(0)
                        i+=1
                    sell_index=i+1
                else:
                    sell_index+=1
                    res_sell.append(0)
            for i in range(len(BUY)):
                if res_buy[i]:
                    res.append(res_buy[i])
                elif res_sell[i]:
                    res.append(res_sell[i])
                else:
                    res.append(0)
        return tdx(res)

def ABS(x: list): 
        return tdx(x).__abs__()

def MAX(x:list,other:Union[list,float,int]):
    return tdx(x).__max__(other)

def MIN(x:list,other:Union[list,float,int]):
    return tdx(x).__min__(other)

def IF(x:list,t_arg: Union[list,float,int],f_arg: Union[list,float,int]):
    return tdx(x).__if__(t_arg,f_arg)

def SMA(x:list,N:int,M:int):
    return tdx(x).__sma__(N,M)

def REF(x:list,window:int):
    return tdx(x).__ref__(window)

def MA(x:list,window:int):
    return tdx(x).__ma__(window)

def LLV(x:list,window:int):
    return tdx(x).__llv__(window)

def HHV(x:list,window:int):
    return tdx(x).__hhv__(window)

def EMA(x:list,window:int):
    return tdx(x).__ema__(window)

def EXPMA(x:list,window:int):
    return tdx(x).__expma__(window)

def CROSS(x:Union[list,float,int],other:Union[list,float,int]):
    if isinstance(x,(float,int)) and isinstance(other,list):
        return tdx(other).__downcross__(x)
    elif isinstance(x,(float,int)) and isinstance(other,(float,int)):
        return tdx([x==other])
    elif isinstance(x,list):
        return tdx(x).__cross__(other)

def STD(x:list,N:int):
    return tdx(x).__std__(N)

def TFILTER(BUY:list,SELL:list,N:int4):
    return tdx(BUY).__tfilter__(SELL,N)


if __name__=='__main__':
    buy =tdx([-5,1,-1])
    sell=tdx([1,0,1])
    tt=tdx([-1,1,2])
    print(buy+sell-sell)
    print(buy*sell*2*sell)
    print(3*buy)
    print(buy*3)
    print(buy/sell)
    print(buy>sell)
    print(buy>=sell)
    print(buy<sell)
    print(buy<=sell)
    print(buy!=sell)
    print(buy>0)
    print(buy>=0)
    print(buy<=0)
    print(buy<0)
    print(buy!=0)
    print(ABS(tt))
    print(MAX(buy,sell))
    print(MIN(buy,sell))
    print(SMA([],4,1))
    print(REF(buy,1))
    print(MA(buy,2))
    print(LLV(buy,2))
    print(HHV(buy,2))
    print(EMA([],2))
    print(EXPMA(buy,2))
    print(CROSS(2,0))
    print(CROSS(0,buy))
    print(CROSS(buy,sell))
    print(STD(buy,2))
    print(TFILTER(buy,sell,2))
    